from protolib import fix_index

import copy

class UndoBuffer:
    def __init__(self):
        self.buffer = []
        self.current_state = None

    def __str__(self):
        return f'<ObjectBuffer {len(self.buffer)}>'

    def add_state(self, state):
        if self.current_state is not None:
            self.buffer.append(self.current_state)
        self.current_state = state

    def undo(self):
        if len(self.buffer) > 0:
            self.current_state = self.buffer.pop()
        return self.current_state

class UndoRedoBuffer: # It's behaviour is a bit weird, but it works good enough
    def __init__(self, initial_state):
        self.buffer = [initial_state]
        self.state_pointer = 0

    def __str__(self):
        return f'<ObjectBuffer2 {len(self.buffer)}>'

    def add_state(self, state):
        if self.state_pointer < len(self.buffer) - 1:
            # If adding a new state, discard any states ahead in the buffer
            self.buffer = self.buffer[: self.state_pointer + 1]
        self.buffer.append(copy.deepcopy(state))
        self.state_pointer = len(self.buffer) - 1

    def get_state(self):
        return self.buffer[self.state_pointer]

    def undo(self):
        if self.state_pointer > 0:
            self.state_pointer -= 1
            self.fix_state_pointer()

    def redo(self):
        if self.state_pointer < len(self.buffer) - 1:
            self.state_pointer += 1
            self.fix_state_pointer()

    def fix_state_pointer(self):
        self.state_pointer = fix_index(self.buffer, self.state_pointer)
        if self.state_pointer == len(self.buffer):
            self.buffer.append(copy.deepcopy(self.buffer[-1]))
